#include "kernel_operator.h"
using namespace AscendC;

template <typename U, typename T>
class KernelInt32
{
    const uint32_t BUFFER_NUM = 2;

public:
    __aicore__ inline KernelInt32() {}
    __aicore__ inline void Init(TPipe *pipeIn, GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t coreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t tailDataNum)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->coreDataNum = coreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = tailDataNum;

        conditionGm.SetGlobalBuffer((__gm__ U *)condition, this->coreDataNum);
        x1Gm.SetGlobalBuffer((__gm__ T *)x1, this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ T *)y, this->coreDataNum);

        pipe = pipeIn;
        pipe->InitBuffer(inQueueCondition, BUFFER_NUM, this->tileDataNum * sizeof(U));
        pipe->InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(T));

        pipe->InitBuffer(tmp1, this->tileDataNum * sizeof(uint8_t));
        pipe->InitBuffer(tmp2, this->tileDataNum * sizeof(half));

        pipe->InitBuffer(tmp3, this->tileDataNum * sizeof(float));
        pipe->InitBuffer(tmp4, this->tileDataNum * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (uint32_t i = 0; i < loopCount; i++)
        {
            if (i == loopCount - 1)
            {
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t progress)
    {
        LocalTensor<U> cdLocal = inQueueCondition.AllocTensor<U>();
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopy(cdLocal, conditionGm[progress * this->tileDataNum], (this->processDataNum + 31) / 32 * 32);
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
        inQueueCondition.EnQue(cdLocal);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    __aicore__ inline void Compute(uint32_t progress)
    {
        LocalTensor<U> cdLocal = inQueueCondition.DeQue<U>();
        LocalTensor<T> x1Local = inQueueX1.DeQue<T>();
        LocalTensor<T> x2Local = inQueueX2.DeQue<T>();
        LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();

        auto selMask = tmp1.Get<uint8_t>();
        auto p1 = tmp2.Get<half>();
        Cast(p1, cdLocal, RoundMode::CAST_NONE, this->processDataNum);
        CompareScalar(selMask, p1, static_cast<half>(1), CMPMODE::EQ, (this->processDataNum + 127) / 128 * 128);

        auto t1 = tmp3.Get<float>();
        auto t2 = tmp4.Get<float>();
        Cast(t1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
        Cast(t2, x2Local, RoundMode::CAST_NONE, this->processDataNum);

        Select(t1, selMask, t1, t2, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
        Cast(yLocal, t1, RoundMode::CAST_ROUND, this->processDataNum);

        outQueueY.EnQue(yLocal);
        inQueueCondition.FreeTensor(cdLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }
    __aicore__ inline void CopyOut(uint32_t progress)
    {
        LocalTensor<T> yLocal = outQueueY.DeQue<T>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, 1> inQueueCondition, inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, 1> outQueueY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    TBuf<QuePosition::VECCALC> tmp3, tmp4;

    GlobalTensor<U> conditionGm;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};
